package org.xteam.sled.solver;

import java.util.ArrayList;
import java.util.Collections;
import java.util.Comparator;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;

import org.xteam.sled.semantic.Range;
import org.xteam.sled.semantic.exp.DefaultExpVisitor;
import org.xteam.sled.semantic.exp.Exp;
import org.xteam.sled.semantic.exp.ExpAdd;
import org.xteam.sled.semantic.exp.ExpCondition;
import org.xteam.sled.semantic.exp.ExpConst;
import org.xteam.sled.semantic.exp.ExpDiv;
import org.xteam.sled.semantic.exp.ExpForce;
import org.xteam.sled.semantic.exp.ExpMod;
import org.xteam.sled.semantic.exp.ExpMul;
import org.xteam.sled.semantic.exp.ExpNarrows;
import org.xteam.sled.semantic.exp.ExpNarrowu;
import org.xteam.sled.semantic.exp.ExpRewriter;
import org.xteam.sled.semantic.exp.ExpShift;
import org.xteam.sled.semantic.exp.ExpSlice;
import org.xteam.sled.semantic.exp.ExpVar;
import org.xteam.sled.semantic.exp.ExpWiden;


public class Balancer {

	private FreshState freshState;
	private Map<StringIntKey, ExpWiden> widens = new HashMap<StringIntKey, ExpWiden>();
	private Map<StringIntKey, ExpNarrows> narrowss = new HashMap<StringIntKey, ExpNarrows>();
	private Map<StringIntKey, ExpNarrowu> narrowus = new HashMap<StringIntKey, ExpNarrowu>();
	private Map<StringIntKey, ExpDiv> divs = new HashMap<StringIntKey, ExpDiv>();
	private Map<StringIntKey, ExpMod> mods = new HashMap<StringIntKey, ExpMod>();
	private Map<String, ExpForce> forces = new HashMap<String, ExpForce>();
	private Map<String, List<Range>> slices = new HashMap<String, List<Range>>();
	
	private List<Balance> balances = new ArrayList<Balance>();
	private Map<Exp, String> balmap = new HashMap<Exp, String>();
	private List<ExpCondition> newEquations = new ArrayList<ExpCondition>();
	private boolean debug;;

	public Balancer(FreshState freshState, boolean debug) {
		this.freshState = freshState;
		this.debug = debug;
	}

	public List<Balance> balanceEqns(List<ExpCondition> equations, List<ExpCondition> constraints, List<ExpCondition> eqns) {
		if (equations.isEmpty())
			return new ArrayList<Balance>();
		debug("=== Balance " + equations);
		InformationCollector collector = new InformationCollector();
		for (ExpCondition equation : equations) {
			if (equation.condition() == ExpCondition.EQ) {
				equation.visit(collector);
				eqns.add(equation);
			} else {
				constraints.add(equation);
			}
		}
		Set<StringPair> nwPairs = new HashSet<StringPair>();
		createWidenBalance(nwPairs);
		createNarrowsBalance(nwPairs);
		createNarrowuBalance();
		createDivModBalance();
		createForceBalance();
		createSliceBalance();
		
		// rewrite equations using balmap
		EquationSubstitution subs = new EquationSubstitution();
		for (ExpCondition eqn : eqns) {
			newEquations.add((ExpCondition)subs.rewrite(eqn));
		}
		eqns.clear();
		eqns.addAll(newEquations);
		debug("equations=" + eqns);
		debug("balances=" + balances);
		return balances;
	}

	private void debug(String msg) {
		if (debug)
			System.out.println(msg);
	}

	private void createWidenBalance(Set<StringPair> nwPairs) {
		for (StringIntKey key : widens.keySet()) {
			ExpWiden outere = widens.get(key);
			String innerv = key.name();
			String outerv = freshState.varFor(outere);
	        if (! nwPairs.contains(new StringPair(innerv, outerv))) {
	        	nwPairs.add(new StringPair(innerv, outerv));
	        	balances.add(new Balance(
	        			new BalanceItem(outerv, new ExpWiden(new ExpVar(innerv), key.value())),
	        			new BalanceItem(innerv, new ExpNarrows(new ExpVar(outerv), key.value()))));
	        }
	        addEquation(new ExpVar(innerv), freshState.var2exp(innerv));
	        balmap.put(outere, outerv);
		}
	}
	
	private void createNarrowsBalance(Set<StringPair> nwPairs) {
		for (StringIntKey key : narrowss.keySet()) {
			ExpNarrows outere = narrowss.get(key);
			String innerv = key.name();
			String outerv = freshState.varFor(outere);
	        if (! nwPairs.contains(new StringPair(outerv, innerv))) {
	        	nwPairs.add(new StringPair(outerv, innerv));
	        	balances.add(new Balance(
	        			new BalanceItem(innerv, new ExpWiden(new ExpVar(outerv), key.value())),
	        			new BalanceItem(outerv, new ExpNarrows(new ExpVar(innerv), key.value()))));
	        }
	        addEquation(new ExpVar(innerv), freshState.var2exp(innerv));
	        balmap.put(outere, outerv);
		}
	}
	
	private void createNarrowuBalance() {
		for (StringIntKey key : narrowus.keySet()) {
			ExpNarrowu e = narrowus.get(key);
			String w = key.name();
			String n = freshState.varFor(e);
			balances.add(new Balance(
					new BalanceItem(w, new ExpVar(n)),
					new BalanceItem(n, new ExpNarrowu(new ExpVar(w), key.value()))));
	        addEquation(new ExpVar(w), freshState.var2exp(w));
	        balmap.put(e, n);
		}
	}
	
	private void createDivModBalance() {
		for (StringIntKey key : divs.keySet()) {
			ExpMod modulus = mods.get(key);
			if (modulus != null) {
				Exp quotient = divs.get(key);

				String q = freshState.varFor(quotient);
				String m = freshState.varFor(modulus);
				Balance balance = new Balance();
				balance.addLeft(new BalanceItem(q, new ExpDiv(key.value(), new ExpVar(key.name()))));
				balance.addLeft(new BalanceItem(m, new ExpMod(key.value(), new ExpVar(key.name()))));
				balance.addRight(new BalanceItem(key.name(),
						new ExpAdd(new ExpMul(key.value(), new ExpVar(q)), new ExpVar(m))));
				balances.add(balance);
				addEquation(new ExpVar(key.name()), freshState
						.var2exp(key.name()));
				balmap.put(quotient, q);
				balmap.put(modulus, m);
			}
		}
	}
	
	private void createForceBalance() {
		for (String r : forces.keySet()) {
			ExpForce f = forces.get(r);
			String i = freshState.varFor(f);
			balances.add(new Balance(
					new BalanceItem(r, new ExpVar(i)),
					new BalanceItem(i, new ExpVar(r))));
			addEquation(new ExpVar(r), freshState.var2exp(r));
			balmap.put(f, i);
		}
	}
	
	private static class RangeComparator implements Comparator<Range> {

		@Override
		public int compare(Range o1, Range o2) {
			if (o1.low() == o2.low())
				return o2.width() - o1.width();
			return o1.low() - o2.low();
		}
		
	}
	
	private void createSliceBalance() {
		for (String v : slices.keySet()) {
			List<Range> sr = slices.get(v);
			Collections.sort(sr, new RangeComparator());
			checkOverlap(sr);
			Exp e = freshState.var2exp(v);
			int bitSize = 32;
			sliceTree(v, e, 0, bitSize, sr);
			addEquation(new ExpVar(v), e);
		}
	}
	
	private List<Range> sliceTree(String w, Exp e, int low, int high, List<Range> slices) {
		if (! slices.isEmpty() && slices.get(0).low() < high)
			return nextSlice(w, e, 0, high, slices, new ArrayList<BalanceItem>(), new ExpConst(0));
		return slices;
	}

	private List<Range> nextSlice(String w, Exp e, int low, int high,
			List<Range> slices, List<BalanceItem> left, Exp sum) {

		if (low == high) {
			if (!slices.isEmpty() && slices.get(0).low() < high)
				throw new RuntimeException("slice consumption");
			Balance balance = new Balance();
			balance.addAllLeft(left);
			balance.addRight(new BalanceItem(w, sum));
			for (BalanceItem item : left) {
				balmap.put(item.getExpr(), item.variable());
			}
			balances.add(balance);
			return slices;
		}
		if (!slices.isEmpty()) {
			Range r = slices.get(0);
			if (r.low() != low) {
				int index = 0;
				for (Range range : slices) {
					if (range.low() >= high) {
						return slices.subList(index, slices.size());
					}
					++index;
				}
				return new ArrayList<Range>();
			} else {
				ExpSlice slice = new ExpSlice(e, r.low(), r.width());
				String s = freshState.varFor(slice);
				List<Range> rest = sliceTree(s, e, low, low + r.width(),
						slices.subList(1, slices.size()));
				left.add(new BalanceItem(s, slice));
				sum = new ExpAdd(sum, new ExpShift(new ExpVar(s), low));
				return nextSlice(w, e, low + r.width(), high, rest, left, sum);
			}
		} else {
			return slices;
		}
	}

	private void checkOverlap(List<Range> sr) {
		for (int i = 0; i < sr.size(); ++i) {
			Range ri = sr.get(i);
			for (int j = i+1; j < sr.size(); ++j) {
				Range rj = sr.get(j);
				if (rj.low() < ri.high() && rj.high() > ri.low()) {
					throw new RuntimeException("overlapping slices");
				}
			}
		}
	}

	private void addEquation(Exp left, Exp right) {
		if (! left.equals(right))
			newEquations .add(new ExpCondition(ExpCondition.EQ, left, right));
	}

	private class InformationCollector extends DefaultExpVisitor {

		public void visitWiden(ExpWiden expWiden) {
			String v = freshState.varFor(expWiden.getExpr());
			widens.put(new StringIntKey(v, expWiden.width()), expWiden);
			super.visitWiden(expWiden);
		}
		
		public void visitNarrows(ExpNarrows expNarrows) {
			String v = freshState.varFor(expNarrows.getExpr());
			narrowss.put(new StringIntKey(v, expNarrows.width()), expNarrows);
			super.visitNarrows(expNarrows);
		}
		
		public void visitNarrowu(ExpNarrowu expNarrowu) {
			String v = freshState.varFor(expNarrowu.getExpr());
			narrowus.put(new StringIntKey(v, expNarrowu.width()), expNarrowu);
			super.visitNarrowu(expNarrowu);
		}
		
		public void visitDiv(ExpDiv expDiv) {
			String v = freshState.varFor(expDiv.getExpr());
			divs.put(new StringIntKey(v, expDiv.factor()), expDiv);
			super.visitDiv(expDiv);
		}
		
		public void visitMod(ExpMod expMod) {
			String v = freshState.varFor(expMod.getExpr());
			mods.put(new StringIntKey(v, expMod.factor()), expMod);
			super.visitMod(expMod);
		}
		
		public void visitSlice(ExpSlice expSlice) {
			String v = freshState.varFor(expSlice.getExpr());
			List<Range> set = slices.get(v);
			if (set == null) {
				slices.put(v, set = new ArrayList<Range>());
			}
			set.add(new Range(expSlice.low(), expSlice.low() + expSlice.width()));
			super.visitSlice(expSlice);
		}
		
		public void visitForce(ExpForce expForce) {
			String v = freshState.varFor(expForce.getExpr());
			forces.put(v, expForce);
			super.visitForce(expForce);
		}
	}
	
	private class EquationSubstitution extends ExpRewriter {
		
		public void visitWiden(ExpWiden expWiden) {
			stack.push(get(expWiden));
		}
		
		public void visitNarrows(ExpNarrows expNarrows) {
			stack.push(get(expNarrows));
		}
		
		public void visitNarrowu(ExpNarrowu expNarrowu) {
			stack.push(get(expNarrowu));
		}
		
		public void visitDiv(ExpDiv expDiv) {
			stack.push(get(expDiv));
		}
		
		public void visitMod(ExpMod expMod) {
			stack.push(get(expMod));
		}
		
		public void visitSlice(ExpSlice expSlice) {
			stack.push(get(expSlice));
		}
		
		private Exp get(Exp exp) {
			if (balmap.containsKey(exp))
				return new ExpVar(balmap.get(exp));
			return exp;
		}
		
	}
	
	private static class StringIntKey {

		private String str;
		private int value;

		public StringIntKey(String str, int value) {
			this.str = str;
			this.value = value;
		}
		
		public int value() {
			return value;
		}

		public String name() {
			return str;
		}

		public int hashCode() {
			return str.hashCode() ^ (value * 37);
		}
		
		public boolean equals(Object o) {
			return ((StringIntKey) o).str.equals(str) && ((StringIntKey) o).value == value;
		}
		
		public String toString() {
			return "{" + str + "," + value + "}";
		}
		
	}
	
	private static class StringPair {

		private String first;
		private String second;

		public StringPair(String first, String second) {
			this.first = first;
			this.second = second;
		}
		
		public int hashCode() {
			return first.hashCode() ^ second.hashCode();
		}
		
		public boolean equals(Object o) {
			return ((StringPair) o).first.equals(first) && ((StringPair) o).second.equals(second);
		}
		
		public String toString() {
			return "{" + first + "," + second + "}";
		}
		
	}

}
